import pandas as pd
import math
import sharppy
import sharppy.sharptab.profile as profile
import sharppy.sharptab.interp as interp
import sharppy.sharptab.winds as winds
import sharppy.sharptab.utils as utils
import sharppy.sharptab.params as params
import sharppy.sharptab.thermo as thermo
import io
import numpy as np

def scp_term(mucape, srh, ebwd):
    if ebwd > 20:
        ebwd = 20.
    elif ebwd < 10:
        ebwd = 0.
     
    muCAPE_term = mucape / 1000.
    esrh_term = srh / 50.
    ebwd_term = ebwd / 20.
    
    scp = muCAPE_term * esrh_term * ebwd_term
    return scp,muCAPE_term,esrh_term,ebwd_term

def stp_fixed_term(sbcape, sblcl, srh01, bwd6):
    if sblcl < 1000.:
        lcl_term = 1.0
    elif sblcl > 2000.:
        lcl_term = 0.0
    else:
        lcl_term = ((2000.-sblcl)/1000.)

    if bwd6 > 30.:
        bwd6 = 30
    elif bwd6 < 12.5:
        bwd6 = 0.0
    
    bwd6_term = bwd6 / 20.

    cape_term = sbcape / 1500.
    srh_term = srh01 / 150.

    stp_fixed = cape_term * lcl_term * srh_term * bwd6_term
   
    return stp_fixed,cape_term,lcl_term,srh_term,bwd6_term
    
def stp_eff_term(mlcape, esrh, ebwd, mllcl, mlcinh):
    cape_term = mlcape / 1500.
    eshr_term = esrh / 150.
    
    if ebwd < 12.5:
        ebwd_term = 0.
    elif ebwd > 30.:
        ebwd_term = 1.5
    else:
        ebwd_term  = ebwd / 20.

    if mllcl < 1000.:
        lcl_term = 1.0
    elif mllcl > 2000.:
        lcl_term = 0.0
    else:
        lcl_term = ((2000. - mllcl) / 1000.)

    if mlcinh > -50:
        cinh_term = 1.0
    elif mlcinh < -200:
        cinh_term = 0
    else:
        cinh_term = ((mlcinh + 200.) / 150.)

    stp_eff = cape_term * eshr_term * ebwd_term * lcl_term * cinh_term
    return stp_eff,cape_term,eshr_term,ebwd_term,lcl_term,cinh_term

def stp_300_term(mucape, srh, bwd, mulcl, mucinh):
    cape_term = mucape / 2000.

    srh_term = srh / 12
    if bwd > 2.85:
        bwd_term = 1.5
    else:
        bwd_term = bwd / 1.9

    if mulcl < 1000.:
        lcl_term = 1.0
    elif mulcl > 1600.:
        lcl_term = 0.0
    else:
        lcl_term = ((1600. - mulcl) / 600.)

    if mucinh > -50:
        cinh_term = 1.0
    elif mucinh < -200:
        cinh_term = 0
    else:
        cinh_term = ((mucinh + 200.) / 150.)

    stp_300 = cape_term * srh_term * bwd_term * lcl_term * cinh_term
    return stp_300, cape_term, srh_term, bwd_term, lcl_term, cinh_term


params=['0-300 m Shear(m/s)','0-1 km Shear(m/s)','0-6 km Shear(m/s)','0-300 m SRH(m2/s2)','effective SRH(m2/s2)','0-1 km SRH(m2/s2)','0-3 km SRH(m2/s2)',
        '4-6km SR wind(m/s)','SBCAPE(J/kg)','MLCAPE(J/kg)','MUCAPE(J/kg)','SWEAT', 'SBCIN(J/kg)', 'MLCIN(J/kg)', 'MUCIN(J/kg)','SBLCL(m AGL)','MLLCL(m AGL)',
        'MULCL(m AGL)','K_index','LI_SFC','850-500hPaLR(C/km)','700-500hPaLR(C/km)', 'lowRH(%)','midRH(%)','0.5-6km BRN shear(m2/s2)','0.5-6km BRN','0-1km EHI',
        '0-3km EHI', '0-4km mean Shear(/s)','VGP','ebwd','SCP','SCP_mucape_term','SCP_esrh_term','SCP_ebwd_term', 'STP(fix)','STP(fix)_cape_term',
        'STP(fix)_lcl_term','STP(fix)_srh_term','STP(fix)_bwd6_term','STP(eff)','STP(eff)_cape_term','STP(eff)_eshr_term','STP(eff)_ebwd_term','STP(eff)_lcl_term',
        'STP(eff)_cinh_term','STP300cn','STP300cn_cape_term','STP300cn_shr300_term', 'STP300cn_shear300_term','STP300cn_lcl_term','STP300cn_cinh_term']
df = pd.DataFrame(columns=params)
f=open("locnontor.txt","r").readlines()
for i in range(1,144):
    print(i)
    a=str(i)    
    fields=f[i-1].split(' ')
    k=int(fields[0])
    lat=str(fields[2])
    lon=str(fields[3])
    ftime=str(fields[1])
    fname="NToS_5/"+a+'_'+ftime+".txt"
    spc_file=open(fname,"r").read()
    import sharppy
    import sharppy.sharptab.profile as profile
    import sharppy.sharptab.interp as interp
    import sharppy.sharptab.winds as winds
    import sharppy.sharptab.utils as utils
    import sharppy.sharptab.params as params
    import sharppy.sharptab.thermo as thermo
    def parseSPC(spc_file):
        data = np.array([l.strip() for l in spc_file.split('\n')])

        title_idx = np.where( data == '%TITLE%')[0][0]
        start_idx = np.where( data == '%RAW%' )[0] +1
        finish_idx = np.where( data == '%END%')[0]

        data_header = data[title_idx + 1].split()
        location = data_header[0]
        time = data_header[1][:11]

        full_data = '\n'.join(data[start_idx[0] : finish_idx[0]][:])
        sound_data =io.StringIO( full_data )

        p, h, T, Td, wdir, wspd = np.genfromtxt( sound_data, delimiter=',', comments="%", unpack=True )
    
        return p, h, T, Td, wdir, wspd

    pres, hght, tmpc, dwpc, wdir, wspd = parseSPC(spc_file)

    prof = profile.create_profile(profile='default', pres=pres, hght=hght, tmpc=tmpc, \
                                        dwpc=dwpc, wspd=wspd, wdir=wdir, missing=-9999, strictQC=True)

    sfcpcl = params.parcelx( prof, flag=1 )
    fcstpcl = params.parcelx( prof, flag=2 )
    mupcl = params.parcelx( prof, flag=3 )
    mlpcl = params.parcelx( prof, flag=4 )

    sfc = prof.pres[prof.sfc]

    p3km = interp.pres(prof, interp.to_msl(prof, 3000.))
    p6km = interp.pres(prof, interp.to_msl(prof, 6000.))
    p1km = interp.pres(prof, interp.to_msl(prof, 1000.))
    p4km = interp.pres(prof, interp.to_msl(prof, 4000.))
    p300m = interp.pres(prof, interp.to_msl(prof, 300.))
    p4kmagl = interp.pres(prof, 4000.)
    p6kmagl = interp.pres(prof, 6000.)
    pa150hPa=sfc-150
    pa300hPa=sfc-300
    mean_3km = winds.mean_wind(prof, pbot=sfc, ptop=p3km)
    sfc_6km_shear = winds.wind_shear(prof, pbot=sfc, ptop=p6km)
    sfc_4km_shear = winds.wind_shear(prof, pbot=sfc, ptop=p4km)
    sfc_3km_shear = winds.wind_shear(prof, pbot=sfc, ptop=p3km)
    sfc_1km_shear = winds.wind_shear(prof, pbot=sfc, ptop=p1km)
    sfc_300m_shear = winds.wind_shear(prof, pbot=sfc, ptop=p300m)
    srwind = params.bunkers_storm_motion(prof)
    srh3km = winds.helicity(prof, 0, 3000., stu = srwind[0], stv = srwind[1])
    srh1km = winds.helicity(prof, 0, 1000., stu = srwind[0], stv = srwind[1])
    srh300m = winds.helicity(prof, 0, 300., stu=srwind[0], stv=srwind[1])
    ship = params.ship(prof)
    eff_inflow = params.effective_inflow_layer(prof)
    ebot_hght = interp.to_agl(prof, interp.hght(prof, eff_inflow[0]))
    etop_hght = interp.to_agl(prof, interp.hght(prof, eff_inflow[1]))
    effective_srh = winds.helicity(prof, ebot_hght, etop_hght, stu = srwind[0], stv = srwind[1])
    el = mupcl.elhght
    de = (el - ebot_hght) / 2
    dep = interp.pres(prof, interp.to_msl(prof, de))
    ebwd = winds.wind_shear(prof, pbot=eff_inflow[0], ptop=dep)
    ebwspd = utils.mag(ebwd[0], ebwd[1]) * 0.5144
    shear0_3 = utils.comp2vec(sfc_300m_shear[0], sfc_300m_shear[1])[1] * 0.5144
    shear0_1=utils.comp2vec(sfc_1km_shear[0], sfc_1km_shear[1])[1]*0.5144
    shear0_6=utils.comp2vec(sfc_6km_shear[0], sfc_6km_shear[1])[1]*0.5144
    mean_shear0_4=utils.comp2vec(sfc_4km_shear[0], sfc_4km_shear[1])[1]*0.5144/4000
    UV4_6km_SRwind=winds.sr_wind(prof,pbot=p4kmagl, ptop=p6kmagl, stu=srwind[0], stv=srwind[1], dp=-1) 
    L4_6km_SRwind=utils.comp2vec(UV4_6km_SRwind[0], UV4_6km_SRwind[1])[1]*0.5144
    sweat=params.sweat(prof)
    k_index=params.k_index(prof)
    L850_500mbLR=params.lapse_rate(prof, 850, 500, pres=True)
    L700_500mbLR=params.lapse_rate(prof, 700, 500, pres=True)
    lowRH=params.mean_relh(prof,pbot=sfc,ptop=pa150hPa,dp=-1,exact=False)
    midRH=params.mean_relh(prof,pbot=pa150hPa,ptop=pa300hPa,dp=-1,exact=False)
    EHI1=srh1km[0]*mlpcl.bplus/160000 
    EHI3=srh3km[0]*mlpcl.bplus/160000
    scp = scp_term(mupcl.bplus, effective_srh[0], ebwspd)
    stp_fixed = stp_fixed_term(sfcpcl.bplus, sfcpcl.lclhght, srh1km[0], utils.comp2vec(sfc_6km_shear[0], sfc_6km_shear[1])[1]*0.5144)
    stp_eff = stp_eff_term(mlpcl.bplus, effective_srh[0], ebwspd, mlpcl.lclhght, mlpcl.bminus)
    stp_300 = stp_300_term(mupcl.bplus, srh300m[0], shear0_3, mupcl.lclhght, mupcl.bminus)
    VGP=mean_shear0_4*math.sqrt(mlpcl.bplus)

    df.loc[i] = [shear0_3,shear0_1,shear0_6,srh300m[0],effective_srh[0],srh1km[0],srh3km[0],L4_6km_SRwind,sfcpcl.bplus,mlpcl.bplus,mupcl.bplus,
                sweat,sfcpcl.bminus,mlpcl.bminus,mupcl.bminus,sfcpcl.lclhght,mlpcl.lclhght,mupcl.lclhght,k_index,sfcpcl.li5,L850_500mbLR,L700_500mbLR,
                lowRH,midRH,sfcpcl.brnshear,sfcpcl.brn,EHI1,EHI3,mean_shear0_4,VGP,ebwspd,scp[0],scp[1],scp[2],scp[3],stp_fixed[0],stp_fixed[1],
                stp_fixed[2],stp_fixed[3],stp_fixed[4],stp_eff[0],stp_eff[1],stp_eff[2],stp_eff[3],stp_eff[4],stp_eff[5],stp_300[0],stp_300[1],stp_300[2],
                stp_300[3],stp_300[4],stp_300[5]]

df.to_excel("NToS_parameter.xlsx",sheet_name="NToS_parameter")
